import socket
import ipaddress
import requests
import concurrent.futures
from typing import List, Optional, Callable, Set

class NetworkScanner:
    def __init__(self, port: int = 11434, timeout: float = 0.5):
        self.port = port
        self.timeout = timeout
        self.custom_networks = set()  # 存储自定义网段
        
    def add_network(self, network: str) -> bool:
        """添加自定义网段
        Args:
            network: 网段，格式如 '192.168.1.0/24'
        Returns:
            bool: 是否添加成功
        """
        try:
            # 验证网段格式
            ipaddress.IPv4Network(network, strict=False)
            self.custom_networks.add(network)
            return True
        except ValueError:
            return False
    
    def remove_network(self, network: str) -> bool:
        """移除自定义网段"""
        try:
            self.custom_networks.remove(network)
            return True
        except KeyError:
            return False
    
    def clear_networks(self):
        """清空自定义网段"""
        self.custom_networks.clear()
    
    def get_networks(self) -> Set[str]:
        """获取所有自定义网段"""
        return self.custom_networks

    def get_local_networks(self) -> List[str]:
        """获取本机所有网络接口的网段"""
        local_networks = set()
        
        try:
            # 获取所有网络接口
            for interface in socket.if_nameindex():
                try:
                    # 获取接口地址信息
                    addrs = socket.getaddrinfo(socket.gethostname(), None)
                    for addr in addrs:
                        ip = addr[4][0]
                        # 只处理IPv4地址
                        if ':' not in ip and not ip.startswith('127.'):
                            network = ipaddress.IPv4Network(f"{ip}/24", strict=False)
                            local_networks.add(str(network))
                except Exception:
                    continue
        except Exception as e:
            print(f"获取网络接口失败: {e}")
            
        return list(local_networks)

    def scan_port(self, ip: str, progress_callback: Optional[Callable] = None) -> Optional[str]:
        """扫描单个IP端口"""
        try:
            # 检查端口是否开放
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.timeout)
            result = sock.connect_ex((str(ip), self.port))
            sock.close()
            
            if result == 0:
                # 测试 Ollama API
                try:
                    url = f"http://{ip}:{self.port}/api/tags"
                    response = requests.get(url, timeout=self.timeout)
                    if response.status_code == 200:
                        return str(ip)
                except Exception as e:
                    if progress_callback:
                        progress_callback(f"API 连接失败 {ip}: {str(e)}")
            return None
        except Exception as e:
            if progress_callback:
                progress_callback(f"扫描错误 {ip}: {str(e)}")
            return None

    def scan_network(self, 
                    progress_callback: Optional[Callable] = None,
                    progress_bar: Optional[Callable] = None) -> List[str]:
        """扫描网络"""
        found_servers = []
        all_ips = set()
        
        # 获取本机所有网段
        local_networks = self.get_local_networks()
        
        # 添加本地网段
        for network in local_networks:
            network_obj = ipaddress.IPv4Network(network, strict=False)
            all_ips.update(network_obj.hosts())
        
        # 添加自定义网段
        for network in self.custom_networks:
            network_obj = ipaddress.IPv4Network(network, strict=False)
            all_ips.update(network_obj.hosts())
        
        # 转换为列表并排序
        all_ips = sorted(list(all_ips))
        total_ips = len(all_ips)
        
        if progress_callback:
            networks_info = local_networks + list(self.custom_networks)
            progress_callback(f"扫描网段: {', '.join(networks_info)}")
            progress_callback(f"正在扫描 {total_ips} 个IP地址...")
        
        with concurrent.futures.ThreadPoolExecutor(max_workers=50) as executor:
            future_to_ip = {
                executor.submit(self.scan_port, ip, progress_callback): ip 
                for ip in all_ips
            }
            
            completed = 0
            for future in concurrent.futures.as_completed(future_to_ip):
                completed += 1
                if progress_bar:
                    progress_bar(completed / total_ips)
                
                result = future.result()
                if result:
                    found_servers.append(result)
                    if progress_callback:
                        progress_callback(f"发现服务器: {result}")
        
        return found_servers 